#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
用于读取traceSeq生成的h5文件并绘制相关图表的脚本
"""

import argparse
import os
import numpy as np
from pylab import plt
import pandas as pd
from obspy.core.utcdatetime import UTCDateTime


# 导入自定义工具
from utils.h5data import read_h5_data
from utils.plot import plot_traces, plot_raw_data
from utils.loc import load_loc, get_distance
from utils.math import norm
from utils.trace import filtfilt

# 命令行参数解析
parser = argparse.ArgumentParser(description='Plot traceSeq h5 data')
parser.add_argument('-input', default='', help='input h5 file')
parser.add_argument('-figroot', default='figures/debug', help='root to save figs')
args = parser.parse_args()

INPUT_FILE = args.input
FIG_ROOT = args.figroot

datasets, args_in_file = read_h5_data(INPUT_FILE, 
                            ['MARKER','EMARKER'],
                            group_name='metadata',
                            read_attrs=True)
MARKER = datasets[0].decode()
EMARKER = datasets[1].decode()
DATE=args_in_file['date']
fs = args_in_file['fs']
fe = args_in_file['fe']
# 根据日期设置参数
if DATE == '2303':
    s_info = load_loc('./loc/loc_all_from_log_2303.csv')
    START_TIME = '2023-03-10T00:00:00Z'
elif DATE == '254C':
    s_info = load_loc('./loc/254C.csv', lat_key='lat', lon_key='lon')
    START_TIME = '2024-10-08T00:00:00Z'

# 读取h5文件中的数据
print(f'Reading data from {INPUT_FILE}')
data_dict = read_h5_data(INPUT_FILE, keys=['traces', 'te_all', 't', 'ref', 't_shift', 'S', 'R', 'all_groups'])
all_traces_slices, te_all, tn, ref_traces_slices, t_shift, S, R, all_groups = data_dict

# 解码字符串数组
CORR_STATS = [name.decode() for name in R]
S = S.decode()

# 创建输出目录
FIG_ROOT = f'{FIG_ROOT}/6A.{DATE}.{EMARKER}'
if not os.path.exists(FIG_ROOT):
    os.makedirs(FIG_ROOT)

# 绘制每个台站的波形图
print('Plotting traces for each station...')
ns = len(CORR_STATS)
for j, name in enumerate(CORR_STATS):
    if j<10 or j>20:
        continue
    plt.close('all')
    traces_j = all_traces_slices[:, j, :]
    traces_j = filtfilt(traces_j,tn[1]-tn[0],[1.5,2.5], order=4)
    fig, ax = plot_traces(traces_j, tn[:], NORM=True, x=(te_all/3600+8)/24)
    ax.set_title(f'{name}.{EMARKER}')
    ax.set_ylabel('N')
    ax.set_xlim([-2,2])
    fig.tight_layout()
    figname = f'{FIG_ROOT}/{name}.{EMARKER}.png'
    fig.savefig(figname)
    print(f'Saved {figname}')
